from .models.DeepLab_v3 import DeepLab_v3
import torch
import torch.nn as nn
from utils.loss import SegmentationLosses
import os.path as osp
from .data.RoadLineDataset import RoadLineDataset
from torch.utils.data import DataLoader, Dataset
import data.DataTXTGenerator as dtg
import os
from tqdm import tqdm
import numpy as np
from utils.metrics import Evaluator
# from visdom import Visdom
import time
import argparse
from utils.parse_config import *

try:
    import moxing as mox
except:
    print('not use moxing')


class Trainer(object):
    def __init__(self, learning_rate=0.001, cuda=False, n_class=8):
        self.model = DeepLab_v3(num_classes=8)

        # self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.001, momentum=0.5,
        #                                  weight_decay=0, nesterov=True)
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)

        self.criterion = SegmentationLosses(weight=None, cuda=cuda, n_class=n_class).build_loss(mode='ce')
        self.evaluator = Evaluator(n_class)
        self.n_class = n_class
        self.train_local=train_local

        # viz.line([0.], [0], win='train_loss', opts=dict(title='train_loss'))
        # viz.line([0.], [0], win='miou', opts=dict(title='miou'))
        # self.viz = viz

    def train(self, epoch, train_loader, batch_size, device):
        train_loss = 0.0
        for i, training_data in enumerate(train_loader, 0):
            inputs, labels = training_data['image'], training_data['label']
            inputs = inputs.to(device)
            labels = labels.to(device)
            self.model.to(device)
            self.optimizer.zero_grad()
            output = self.model(inputs)
            loss = self.criterion(output, labels)
            loss.backward()
            self.optimizer.step()
            train_loss += loss.item()
        print('---[Epoch: %d]' % (epoch))
        print('---Training Loss: %.3f' % train_loss)

        # self.viz.line([train_loss / batch_size], [epoch], win='train_loss', update='append')

    def validation(self, epoch, validat_loader, batch_size, device):
        validation_loss = 0.0
        batch_total_miou = 0.0
        for i, training_data in enumerate(validat_loader, 0):
            inputs, labels = training_data['image'], training_data['label']
            inputs = inputs.to(device)
            labels = labels.to(device)
            self.model.to(device)
            with torch.no_grad():
                output = self.model(inputs)
            loss = self.criterion(output, labels)
            validation_loss += loss.item()
            pred = output.data.cpu()
            target = labels.cpu()
            pred = torch.argmax(pred, dim=1)
            # Add batch sample into evaluator
            # self.evaluator.add_batch(target, pred)

            #     mIOU Calculation
            miou = self.calculate_miou(pred, target)
            print("--miou:", miou)
            batch_total_miou += miou.item()
        # self.viz.line([batch_total_miou / batch_size], [epoch], win='miou', update='append')
        # time.sleep(0.5)

    def calculate_miou(self, predicts, labels):
        ious = torch.zeros(self.n_class)
        for i in range(self.n_class):
            intersection = torch.sum((predicts == i) * (labels == i))
            area = torch.sum(predicts == i) + torch.sum(labels == i) - intersection
            ious[i] = intersection / (area + 1e-6)
        miou = torch.sum(ious) / (self.n_class + 1)
        return miou


def prepare_data_on_modelarts(pretrained_weights, local_data_root, data_url, train_url):
    """
    将OBS上的数据拷贝到ModelArts中
    """
    # 拷贝预训练参数文件

    # 默认使用ModelArts中的如下两个路径用于存储数据：
    # 0) /cache/model: 如果适用预训练模型，存储从OBS拷贝过来的预训练模型
    # 1）/cache/datasets: 存储从OBS拷贝过来的训练数据
    # 2）/cache/log: 存储训练日志和训练模型，并且在训练结束后，该目录下的内容会被全部拷贝到OBS
    if pretrained_weights:
        _, weights_name = os.path.split(pretrained_weights)
        mox.file.copy(pretrained_weights, os.path.join(local_data_root, 'model/' + weights_name))
        pretrained_weights = os.path.join(local_data_root, 'model/' + weights_name)

    data_local = os.path.join(local_data_root, 'datasets/processed_data')
    if not os.path.exists(data_local):
        data_dir = os.path.join(local_data_root, 'datasets')
        print("data_dir:",data_dir)
        mox.file.copy_parallel(data_url, data_dir)
        os.system('cd %s;unzip processed_data.zip' % data_dir)  # 训练集已提前打包为trainval.zip
        print("data_local:",data_local)
        if os.path.isdir(data_local):
            os.system('cd %s;rm processed_data.zip' % data_dir)
            print('unzip processed_data.zip success, args.data_local is', data_local)
        else:
            raise Exception('unzip processed_data.zip Failed')
    else:
        print('args.data_local: %s is already exist, skip copy' % data_local)

    output_path = os.path.join(local_data_root, 'log/')
    if not os.path.exists(output_path):
        os.mkdir(output_path)

    return pretrained_weights, data_local, output_path


if __name__ == "__main__":
    # python -m visdom.server

    EPOCH = 50
    BATCH_SIZE = 32
    learning_rate = 0.0001
    base_size = 64
    pretrained_weights = ''
    # train_local = True
    train_local = False
    if train_local:
        # data_path = '../RoadLineDataset'
        data_config_path = 'data/custom.data'
        pretrained_weights = ''
        data_local='../RoadLineDataset/processed_data'
        output_path=''
    else:
        # Get data configuration
        data_config_path = 'LaneLineDetection/data/custom.data'
        local_data_root = '/cache/'
        data_url = 'obs://class-1275-42261/Lab-2210/modelarts22905308/datasets/'
        train_url = 'obs:///class-1275-42261/Lab-2210/modelarts22905308/train_output/V0015/'
        pretrained_weights = ''
        pretrained_weights, data_local, output_path = prepare_data_on_modelarts(pretrained_weights, local_data_root,
                                                                                data_url, train_url)
        print("output_path:",output_path)
        print("data_local:", data_local)
    dtg.creatRandomDataTxt(train_local,data_local, 0.8,0.1)
    data_config = parse_data_config(data_config_path)
    train_path = data_config["local_train"] if train_local else data_config["train"]
    valid_path = data_config["local_valid"] if train_local else data_config["valid"]
    train_data = RoadLineDataset(train_path, base_size=base_size)
    valid_data = RoadLineDataset(valid_path, base_size=base_size)
    train_data_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)  # 使用DataLoader加载数据
    valid_data_loader = DataLoader(valid_data, batch_size=BATCH_SIZE, shuffle=True)  # 使用DataLoader加载数据
    # tbar = tqdm(data_loader)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("devide:", device)
    trainer = Trainer(learning_rate=learning_rate, cuda=torch.cuda.is_available())

    for epoch in range(EPOCH):
        trainer.train(epoch, train_data_loader, BATCH_SIZE, device)
        # trainer.validation(epoch, data_loader, BATCH_SIZE, device)
        if epoch % 5 == 0:
            AP = trainer.validation(epoch, valid_data_loader, BATCH_SIZE, device)
            temp_model_name = f"ckpt_%d.pth" % (epoch)
            ckpt_name = os.path.join('models', 'save', temp_model_name)
            if not train_local:
                ckpt_name = os.path.join('LaneLineDetection', ckpt_name)
            torch.save(trainer.model.state_dict(), ckpt_name)
            # save model to obs
            if not train_local:
                mox.file.copy_parallel(ckpt_name, os.path.join(train_url, temp_model_name))

# if not train_local:
#     mox.file.copy_parallel(., os.path.join(opt.train_url, temp_model_name))
